'''
Author: devis.dong
Email: devis.dong@gmail.com
Date: 2021-12-08 15:20:41
LastEditTime: 2022-12-11 02:11:52
LastEditors: devis.dong
Description:
'''

import os
import sys
import math
import torch
import random
import string
import logging
from datetime import datetime


def timenow():
    return datetime.now().strftime('%Y%m%d%H%M%S')

def getLogger(name=None, console_level=logging.DEBUG, file_level=logging.DEBUG, file_log=os.path.join('logs', timenow() + '.log')):
    os.makedirs(os.path.abspath('logs'), exist_ok=True)
    file_log = os.path.join('logs', timenow() + '.log')
    with open(file_log, 'w') as f:
            pass
    logfile_handler = logging.FileHandler(file_log, mode='a', encoding="utf8")
    logfile_handler.setFormatter(logging.Formatter('%(asctime)s [%(levelname)s] TID.%(thread)d %(module)s.%(lineno)d %(name)s:\t%(message)s'))
    logfile_handler.setLevel(file_level)
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(logging.Formatter('[%(asctime)s %(levelname)s] %(message)s',datefmt="%Y/%m/%d %H:%M:%S"))
    console_handler.setLevel(console_level)
    logger = logging.getLogger(name)
    logger.addHandler(logfile_handler)
    logger.addHandler(console_handler)
    logger.setLevel(min(console_level, file_level))
    return logger

logger = getLogger()


def save_checkpoint(model, epoch, ckpt_name, accu_all=0, accu_pos=0, accu_neg=0, accu_cls=0):
    if isinstance(model, torch.nn.DataParallel):
        model_state = model.module.state_dict()
    else:
        model_state = model.state_dict()

    state = {'epoch': epoch, 'model_state': model_state, 'accu_all': accu_all, 'accu_pos': accu_pos, 'accu_neg': accu_neg, 'accu_cls':accu_cls}
    torch.save(state, ckpt_name)

def load_checkpoint(model, filename):
    epoch, accu_all, accu_pos, accu_neg, accu_cls = 1, 0, 0, 0, 0
    if os.path.isfile(filename):
        logger.info("==> Loading from checkpoint %s" % filename)
        checkpoint = torch.load(filename, map_location='cpu')
        print(checkpoint.keys())
        epoch = checkpoint['epoch']
        accu_all = checkpoint['accu_all']
        accu_pos = checkpoint['accu_pos']
        accu_neg = checkpoint['accu_neg']
        accu_cls = checkpoint['accu_cls']
        # accu_cls = checkpoint['metric']

        model_saved = checkpoint['model_state']
        model_dict =  model.state_dict()
        state_dict = {k:v for k,v in model_saved.items() if k in model_dict.keys()}
        # print(state_dict.keys())  # dict_keys(['w', 'conv1.weight', 'conv1.bias', 'conv2.weight', 'conv2.bias'])
        model_dict.update(state_dict)
        model.load_state_dict(model_dict)
        logger.info("==> Done")
    else:
        logger.info("file %s does not exist." % filename)

    return epoch, accu_all, accu_pos, accu_neg, accu_cls

def load_model(filepath):
    if not os.path.isfile(filepath):
        logger.error('Invalid file path.')
        return None

    logger.info("==> Loading model from %s." % filepath)
    model = torch.load(filepath)
    logger.info("==> Load model from %s successfully." % filepath)
    return model

def save_model(model, filepath):
    torch.save(model, filepath)

def tensor2str(feature:torch.tensor, delimiter=','):
  return delimiter.join(map(str, feature.tolist()))

def str2tensor(feature:str, delimiter=','):
  return torch.tensor(list(map(float, feature.split(delimiter))))

def calculate_similarity(feature0:torch.tensor, feature1:torch.tensor):
    euclidean_distance = torch.sum((feature0-feature1)**2, dim=-1)
    simi = torch.exp(-euclidean_distance)
    return simi

def random_str(length=10):
    return ''.join([random.choice(string.ascii_letters) for i in range(length)])

def read_pointfile(txt_file):
    data = [] # [N, 3]
    print('reading file', txt_file, '....')
    with open(txt_file, 'r') as f:
        for line in f.readlines():
            data.append(list(map(float, line.strip().split(','))))
    print('reading file', txt_file, 'done.')
    return data

def read_plainfile(filepath, line_delimiter='\n'):
    data = ""
    print('reading file', filepath, '....')
    with open(filepath, 'r') as f:
        lines = f.readlines()
        data = line_delimiter.join(lines)
    print('reading file', filepath, 'done.')
    return data

def save_train_data(filename, epoch, accu_all_train=0, accu_pos_train=0, accu_neg_train=0,
                            accu_all_valid=0, accu_pos_valid=0, accu_neg_valid=0, accu_cls_train=0, accu_cls_valid=0 ):
    with open(filename, "a+") as f:
        f.write("%d,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f\n"
                % (epoch, accu_all_train, accu_pos_train, accu_neg_train, accu_all_valid, accu_pos_valid, accu_neg_valid, accu_cls_train, accu_cls_valid))


#config_logging(os.path.join('logs/', timenow() + '.log'))